package com.smh.aiagent.rag;

import jakarta.annotation.Resource;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.preretrieval.query.transformation.RewriteQueryTransformer;
import org.springframework.stereotype.Component;

/**
 * @author shiminghui
 * @date 2025/6/3 16:26
 * @description: TODO
 */
@Component
public class MyQueryRewriter {


    @Resource
    private ChatModel dashscopeChatModel;

    public Query rewrite(Query query) {

        ChatClient.Builder builder = ChatClient.builder(dashscopeChatModel);
        RewriteQueryTransformer queryTransformer = RewriteQueryTransformer
                .builder()
                .chatClientBuilder(builder)
                .build();
        return queryTransformer.transform(query);
    }


}
